# Copyright 2021 The IREE Authors
#
# Licensed under the Apache License v2.0 with LLVM Exceptions.
# See https://llvm.org/LICENSE.txt for license information.
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
"""
Sets up a locally symlinked llvm-project directory and configures bazel to
use it.

In its default mode of operation, TensorFlow provides the llvm instance by
downloading a blob from GitHub and expanding it. Run this script to switch
it instead to use the llvm-project from IREE's third_party directory. Aside
from saving download time, this lets you make local edits.
"""

import os
import shutil


def symlink_dir_children(source_dir, dest_dir):
  print(f"Symlink children from {source_dir} -> {dest_dir}")
  os.makedirs(dest_dir, exist_ok=True)
  for child in os.listdir(source_dir):
    dest_path = os.path.join(dest_dir, child)
    if os.path.exists(dest_path):
      os.unlink(dest_path)
    os.symlink(os.path.join(source_dir, child), dest_path)


this_dir = os.path.dirname(os.path.abspath(__file__))
workspace_root = this_dir
for i in range(2):
  workspace_root = os.path.dirname(workspace_root)
repo_root = workspace_root
for i in range(2):
  repo_root = os.path.dirname(repo_root)
print(f"IREE repository root: {repo_root}")

llvm_dir = os.path.join(repo_root, "third_party", "llvm-project")
bazel_llvm_dir = os.path.join(this_dir, ".deps", "llvm-project")
tensorflow_dir = os.path.join(repo_root, "third_party", "tensorflow")

print(f"Symlinking shadow LLVM directory into: {bazel_llvm_dir}")
os.makedirs(bazel_llvm_dir, exist_ok=True)
for name in ["llvm", "mlir"]:
  symlink_dir_children(os.path.join(llvm_dir, name),
                       os.path.join(bazel_llvm_dir, name))

# Now copy some bits and bobs from tensorflow.
with open(os.path.join(bazel_llvm_dir, "WORKSPACE"), "wt") as f:
  f.write("""workspace(name = "llvm-project")""")

shutil.copyfile(
    os.path.join(tensorflow_dir, "third_party", "llvm",
                 "llvm.autogenerated.BUILD"),
    os.path.join(bazel_llvm_dir, "llvm", "BUILD"))
shutil.copyfile(os.path.join(tensorflow_dir, "third_party", "mlir", "BUILD"),
                os.path.join(bazel_llvm_dir, "mlir", "BUILD"))

rcfile = os.path.join(workspace_root, "tf-user.bazelrc")
print(f"Writing rc file: {rcfile}")
with open(rcfile, "wt") as f:
  f.write(f"build --override_repository=llvm-project={bazel_llvm_dir}\n")
